from PIL import Image
import matplotlib
matplotlib.use('TkAgg')
from pandasai import SmartDataframe
from pandasai.callbacks import StdoutCallback
from pandasai.llm import OpenAI
import streamlit as st
import pandas as pd
import os
from pandasai.responses.response_parser import ResponseParser

# 处理数据显示
class StreamlitResponse(ResponseParser):
    def __init__(self, context) -> None:
        super().__init__(context)

    def format_dataframe(self, result):
        st.dataframe(result["value"])
        return  result

    def format_plot(self, result):
        img = Image.open(result["value"])
        st.image(img)
        return result

    def format_other(self, result):
        st.write(result["value"])
        return result

st.set_page_config(page_title="Chatbot For Data", page_icon=":school:", menu_items={
    'Get Help': 'mailto:2301366947@qq.com',
    'Report a bug': "mailto:2301366947@qq.com",
    'About': '您可以提问有关excel和数据库的问题，本程序将给出您答案'
})

# 数据准备
file_formats = {
    "csv": pd.read_csv,
    "xls": pd.read_excel,
    "xlsx": pd.read_excel,
    "xlsm": pd.read_excel,
}
openai_api_key = "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"


# 标定历史记录数据文件名
if "file_names" not in st.session_state:
    st.session_state["file_names"] = set()

# 清除提交
def clear_submit():
    st.session_state["submit"] = False

# 数据缓存
@st.cache_data(ttl="2h")
def load_data(uploaded_file):
    try:
        ext = os.path.splitext(uploaded_file.name)[1][1:].lower()
    except:
        ext = uploaded_file.split(".")[-1]
    if ext in file_formats:
        return file_formats[ext](uploaded_file)
    else:
        st.error(f"Unsupported file format: {ext}")
        return None

# 界面UI
st.write("<center><h2>基于PandasAI的智能知识库<h2></center>", unsafe_allow_html=True)
with st.sidebar:
    st.image('pku_logo.png')
    st.markdown("## 数据加载")
    with st.header("Option"):
        uploaded_file = st.file_uploader(
            '数据加载',
            type=list(file_formats.keys()),
            help="可支持的文件类型为：csv/xls/xlsx/xlsm",
            on_change=clear_submit,
            label_visibility="hidden",
        )
if uploaded_file:
    df = load_data(uploaded_file)
    with st.expander("🔎 部分数据预览"):
        st.write(df.tail(5))
    st.session_state.file_names.add(uploaded_file.name)
st.sidebar.markdown("---")
st.sidebar.markdown("## 历史记录")
option = st.sidebar.selectbox(
    '历史记录',
    st.session_state.file_names,
    index=0,     # 默认选最后提交作为聊天记录
    placeholder="请选择提问文件",
    label_visibility="hidden",
)

if option:
    chatData = str(option.split('.')[0])     # 文件名作为索引
    # 根据文件名存储对话记录
    # 初始化提问
    if chatData not in st.session_state or st.sidebar.button("   清空当前数据提问记录   "):
        st.session_state[chatData] = [{"role": "assistant", "content": "我可以帮到您什么?"}]
    # 加载所有对话
    for msg in st.session_state[chatData]:
        # 处理问题的显示
        if isinstance(msg["content"], str):
            st.chat_message(msg["role"]).write(msg["content"])
        # 处理结果的显示
        elif msg["content"]["type"] == "dataframe":
            st.chat_message(msg["role"]).dataframe(msg["content"]['value'])
        elif msg["content"]["type"] == "plot":
            img = Image.open(msg["content"]['value'])
            st.chat_message(msg["role"]).image(img)
        else:
            st.chat_message(msg["role"]).write(msg["content"]['value'])
    if prompt := st.chat_input(placeholder="请输入问题..."):
        # 加入list
        st.session_state[chatData].append({"role": "user", "content": prompt})
        # 写入对话框
        st.chat_message("user").write(prompt)

        llm = OpenAI(api_token=openai_api_key)
        # 处理df未定义
        sdf = SmartDataframe(df,
                             config={
                                 "llm": llm,
                                 "enable_cache": True,
                                 "response_parser": StreamlitResponse,
                                 "callback": StdoutCallback()
                             })

        sdf.enforce_privacy = True

        with st.chat_message("assistant"):
            with st.spinner("加载中..."):
                # pandasAI进行处理
                response = sdf.chat(prompt)
            # 写入结果
            print("-------------------------")
            print(response)
            st.session_state[chatData].append({"role": "assistant", "content": response})
else:
    prompt = st.chat_input(placeholder="请输入问题...")
    st.chat_message("assistant").write("请在左侧上传要分析的数据集 ~")
